"""
边特征聚合操作

提供与DGL兼容的边特征聚合操作，使用PyG作为后端
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
import sys
import os

# 添加路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(os.path.dirname(current_dir))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

from ops.src.graph.graph import is_dgl_graph


def copy_e_sum(graph, edge_features):
    """
    将边特征聚合到目标节点
    
    参数:
        graph: Graph对象或DGL图
        edge_features: 边特征张量
    
    返回:
        聚合后的节点特征
    """
    print("===================copy_e_sum===================")
    # print(f"Input edge_features shape: {edge_features.shape}")
    # 检查图是否为DGL图
    if is_dgl_graph(graph):
        # DGL图
        import dgl.function as fn
        graph = graph.local_var()
        graph.edata['e'] = edge_features
        graph.update_all(fn.copy_e('e', 'm'), fn.sum('m', 'out'))
        result = graph.ndata['out']
        # print(f"DGL output shape: {result.shape}")
        return result
    else:
        # PyG风格的Graph
        dst = graph.edge_index[1]
        # print(f"Destination nodes (dst) shape: {dst.shape}")
        
        # 确保num_nodes是整数
        if hasattr(graph, 'num_nodes'):
            if callable(graph.num_nodes):
                num_nodes = graph.num_nodes()
            else:
                num_nodes = graph.num_nodes
        else:
            # 如果没有num_nodes属性，则从边索引中推断
            num_nodes = int(dst.max().item() + 1)
        
        # print(f"Number of nodes: {num_nodes}")
            
        dtype = edge_features.dtype
        device = edge_features.device
        
        # 使用高效的PyTorch原生操作替代Python循环
        # 创建输出张量
        if edge_features.dim() == 1:
            result = torch.zeros(int(num_nodes), dtype=dtype, device=device)
        else:
            output_shape = [int(num_nodes)] + [int(s) for s in edge_features.shape[1:]]
            result = torch.zeros(output_shape, dtype=dtype, device=device)
        
        # print(f"Result tensor shape before aggregation: {result.shape}")
        
        # 使用index_add_进行高效聚合（比Python循环快100倍以上）
        result.index_add_(0, dst, edge_features)
        
        # print(f"Result tensor shape after aggregation: {result.shape}")
        # print("=====================copy_e_sum_===========================")
        return result 